#pragma once

#include <thrust/replace.h>

template <typename Container,
          typename T>
struct Replace
{
  Container A, A_copy;
  T old_value, new_value;

  template <typename Range>
  Replace(const Range& X, const T& old_value, const T& new_value)
    : A(X.begin(), X.end()), A_copy(X.begin(), X.end()),
      old_value(old_value), new_value(new_value)
  {}

  void operator()(void)
  {
    thrust::replace(A.begin(), A.end(), old_value, new_value);
  }
  
  void reset(void)
  {
    // restore initial data
    thrust::copy(A_copy.begin(), A_copy.end(), A.begin());
  }
};

template <typename Container1,
          typename Container2,
          typename T>
struct ReplaceCopy
{
  Container1 A;
  Container2 B;
  T old_value, new_value;

  template <typename Range1, typename Range2>
  ReplaceCopy(const Range1& X, const Range2& Y, const T& old_value, const T& new_value)
    : A(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      old_value(old_value), new_value(new_value)
  {}

  void operator()(void)
  {
    thrust::replace_copy(A.begin(), A.end(), B.begin(), old_value, new_value);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container,
          typename Predicate,
          typename T>
struct ReplaceIf
{
  Container A, A_copy;
  Predicate pred;
  T new_value;

  template <typename Range>
  ReplaceIf(const Range& X, Predicate pred, const T& new_value)
    : A(X.begin(), X.end()), A_copy(X.begin(), X.end()),
      pred(pred), new_value(new_value)
  {}

  void operator()(void)
  {
    thrust::replace_if(A.begin(), A.end(), pred, new_value);
  }
  
  void reset(void)
  {
    // restore initial data
    thrust::copy(A_copy.begin(), A_copy.end(), A.begin());
  }
};

template <typename Container1,
          typename Container2,
          typename Predicate,
          typename T>
struct ReplaceIfStencil
{
  Container1 A, A_copy;
  Container2 B;
  Predicate pred;
  T new_value;

  template <typename Range1, typename Range2>
  ReplaceIfStencil(const Range1& X, const Range2& Y, Predicate pred, const T& new_value)
    : A(X.begin(), X.end()), A_copy(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      pred(pred), new_value(new_value)
  {}

  void operator()(void)
  {
    thrust::replace_if(A.begin(), A.end(), B.begin(), pred, new_value);
  }
  
  void reset(void)
  {
    // restore initial data
    thrust::copy(A_copy.begin(), A_copy.end(), A.begin());
  }
};

template <typename Container1,
          typename Container2,
          typename Predicate,
          typename T>
struct ReplaceCopyIf
{
  Container1 A, A_copy; // input
  Container2 B;         // output
  Predicate pred;
  T new_value;

  template <typename Range1, typename Range2>
  ReplaceCopyIf(const Range1& X, const Range2& Y, Predicate pred, const T& new_value)
    : A(X.begin(), X.end()), A_copy(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      pred(pred), new_value(new_value)
  {}

  void operator()(void)
  {
    thrust::replace_copy_if(A.begin(), A.end(), B.begin(), pred, new_value);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

template <typename Container1,
          typename Container2,
          typename Container3,
          typename Predicate,
          typename T>
struct ReplaceCopyIfStencil
{
  Container1 A, A_copy; // input
  Container2 B;         // stencil
  Container3 C;         // output
  Predicate pred;
  T new_value;

  template <typename Range1, typename Range2, typename Range3>
  ReplaceCopyIfStencil(const Range1& X, const Range2& Y, const Range3& Z, Predicate pred, const T& new_value)
    : A(X.begin(), X.end()), A_copy(X.begin(), X.end()),
      B(Y.begin(), Y.end()),
      C(Z.begin(), Z.end()),
      pred(pred), new_value(new_value)
  {}

  void operator()(void)
  {
    thrust::replace_copy_if(A.begin(), A.end(), B.begin(), C.begin(), pred, new_value);
  }
  
  void reset(void)
  {
    // nothing to do
  }
};

